import re
from typing import Tuple, Optional
import os
from biomni.llm import get_llm
from logger_config import get_logger

logger = get_logger(__name__)

def classify_conversation_with_llm(query: str) -> bool:
    """
    使用LLM判断对话是否为简单对话
    
    Args:
        query: 用户查询文本
        
    Returns:
        bool: True表示简单对话，False表示复杂任务
    """
    try:
        # 获取LLM实例，从环境变量读取配置
        api_key = os.getenv('API_KEY', 'EMPTY')
        base_url = os.getenv('BASE_URL')
        model = os.getenv('DEFAULT_MODEL', 'claude-sonnet-4-20250514')
        
        llm = get_llm(
            model=model,
            api_key=api_key,
            base_url=base_url
        )
        
        # 设计分类prompt
        classification_prompt = f"""你是一个生物医学AI助手的对话分类器。请判断用户的查询是"简单对话"还是"复杂任务"。

判断标准：

**简单对话**包括：
- 问候语（你好、hi、hello等）
- 感谢和告别（谢谢、再见、bye等）
- 询问助手功能和能力（"你能做什么"、"你有什么功能"等）
- 基础生物学概念解释（不需要数据分析或计算）
- 简单的事实性问题
- 一般性咨询

**复杂任务**包括：
- 需要数据分析、计算或建模的请求
- CRISPR实验设计
- 基因组学、蛋白质组学分析
- 文献检索和综合
- 实验方案设计
- 生物信息学分析
- 涉及具体研究项目的深度问题
- 需要使用专业工具和数据库的任务

用户查询："{query}"

请只回答"简单对话"或"复杂任务"，不要添加任何解释。"""
        
        # 调用LLM进行分类
        response = llm.invoke(classification_prompt)
        result = response.content.strip() if hasattr(response, 'content') else str(response).strip()
        
        # 解析结果
        is_simple = "简单对话" in result
        
        logger.info(f"LLM分类结果: 查询='{query[:50]}...', 分类='{result}', 简单对话={is_simple}")
        return is_simple
        
    except Exception as e:
        logger.error(f"LLM分类失败，回退到关键词匹配: {e}")
        # 如果LLM调用失败，回退到原有的关键词匹配逻辑
        return is_simple_conversation_keyword_based(query)

def is_simple_conversation_keyword_based(query: str) -> bool:
    """
    基于关键词匹配的对话分类（备用方案）
    
    Args:
        query: 用户查询文本
        
    Returns:
        bool: True表示简单对话，False表示复杂任务
    """
    query_lower = query.lower().strip()
    
    # 问候语和基本交互
    greetings = ['你好', 'hello', 'hi', '嗨', '您好']
    if any(greeting in query_lower for greeting in greetings):
        return True
    
    # 简单问题关键词
    simple_keywords = [
        '是什么', '什么是', '能做什么', '有什么功能', '怎么用',
        '谢谢', '感谢', '再见', 'bye', 'thanks', 'thank you'
    ]
    if any(keyword in query_lower for keyword in simple_keywords):
        return True
    
    # 复杂生物学关键词
    complex_bio_keywords = [
        'crispr', '基因编辑', '测序', 'rna-seq', 'chip-seq', 'atac-seq',
        '蛋白质组', '基因组', '生物信息', '实验设计', '数据分析',
        '建模', '预测', '筛选', '文献检索'
    ]
    if any(keyword in query_lower for keyword in complex_bio_keywords):
        return False
    
    # 复杂动词
    complex_verbs = ['分析', '设计', '构建', '预测', '筛选', '比较', '识别']
    if any(verb in query_lower for verb in complex_verbs):
        return False
    
    # 查询长度判断
    if len(query) > 100:
        return False
    
    # 生物学专业术语
    bio_terms = ['基因', '蛋白质', '细胞', '分子', '通路', '疾病', '治疗']
    bio_term_count = sum(1 for term in bio_terms if term in query_lower)
    if bio_term_count >= 2:
        return False
    
    return True

def is_simple_conversation(query: str) -> bool:
    """
    判断用户查询是否为简单对话
    
    Args:
        query: 用户查询文本
        
    Returns:
        bool: True表示简单对话，False表示复杂任务
    """
    # 使用环境变量控制是否启用LLM分类
    use_llm_classifier = os.environ.get('USE_LLM_CLASSIFIER', 'true').lower() == 'true'
    
    if use_llm_classifier:
        return classify_conversation_with_llm(query)
    else:
        return is_simple_conversation_keyword_based(query)

def get_simple_response(query: str) -> str:
    """
    使用LLM为简单对话生成回复
    
    Args:
        query: 用户查询文本
        
    Returns:
        str: 动态生成的简单回复
    """
    try:
        # 获取LLM实例，从环境变量读取配置
        api_key = os.getenv('API_KEY', 'EMPTY')
        base_url = os.getenv('BASE_URL')
        model = os.getenv('DEFAULT_MODEL', 'claude-sonnet-4-20250514')
        
        llm = get_llm(
            model=model,
            api_key=api_key,
            base_url=base_url
        )
        
        # 设计简单对话回复prompt
        response_prompt = f"""你是biomni，一个生物医学AI助手。请根据用户的简单对话查询，生成合适的回复。

你的身份和能力：
- 你是biomni，专门的生物医学人工智能助手
- 你专注于基因组学、蛋白质组学、生物信息学和临床医学
- 你可以回答生物学问题、解释概念、分析数据、设计实验
- 你可以自主执行复杂的生物医学工作流程
- 你可以使用专门的生物医学工具和数据库
- 你在25+生物医学领域工作

你的主要能力包括：
**直接协助：**
• 回答有关生物学、基因组学、蛋白质组学和临床医学的事实问题
• 解释生物学概念、途径和机制
• 提供有关基因、蛋白质、疾病和治疗的信息
• 讨论实验技术和方案

**计算和研究任务：**
• 分析基因组、蛋白质组或临床数据集
• 设计CRISPR实验和筛选
• 进行生物信息学分析
• 检索和综合科学文献
• 生成研究方案
• 根据生物医学数据构建预测模型
• 识别疾病相关基因和途径
• 比较治疗方案或实验方法

**具体领域能力：**
• 基础生物学问题：解释基因、蛋白质、细胞生物学概念，讨论分子生物学技术
• 基因组学分析：分析测序数据、基因表达差异分析、基因功能注释
• 蛋白质组学：质谱数据分析、蛋白质相互作用网络分析
• CRISPR实验设计：设计sgRNA和实验方案、规划基因敲除/敲入实验
• 临床医学应用：疾病相关基因识别、药物靶点分析、临床数据分析
• 文献检索与协议生成：搜索科研文献、生成实验协议、比较研究方法

用户查询："{query}"

请生成一个自然、友好、专业的回复。如果是问候，要介绍你的身份和能力。如果询问功能，要详细说明你能做什么。保持回复简洁但信息丰富。"""
        
        # 调用LLM生成回复
        response = llm.invoke(response_prompt)
        result = response.content.strip() if hasattr(response, 'content') else str(response).strip()
        
        logger.info(f"LLM生成简单回复: 查询='{query[:50]}...', 回复长度={len(result)}")
        return result
        
    except Exception as e:
        logger.error(f"LLM生成简单回复失败，使用默认回复: {e}")
        # 如果LLM调用失败，返回默认回复
        return "你好！我是biomni，一个生物医学AI助手。我可以帮助您解答生物学问题、分析数据或提供研究建议。请告诉我您需要什么帮助？"